from model.resnet import ResNet
from model.senet import se_resnet
from model.senet.se_inception import se_inception_v3
from model.efficientnet import Efficientnet
from model.resnet_double import ResNet_double
from model.double.resnet import resnet18,resnet50,resnet34
from model.double.senet import se_resnet50,se_resnet18,se_resnet56
from model.resnet3d import  resnet,senet
from model.efficientnetv2.model import efficientnetv2_s 
from vit_pytorch.efficient import ViT
from linformer import Linformer

import torchvision
def get_single_model(modelname,num_classes=2,in_channel=3):
    if modelname =='resnet18':
        model = ResNet(model =torchvision.models.resnet18(pretrained=True),num_classes=num_classes,in_channel=in_channel)
    elif modelname =='resnet34':
        model = ResNet(model =torchvision.models.resnet34(pretrained=True),num_classes=num_classes,in_channel=in_channel)
    elif modelname =='resnet50':
        model = ResNet(model =torchvision.models.resnet50(pretrained=True),num_classes=num_classes,in_channel=in_channel)
    elif modelname =='resnext50_32x4d':
        model = ResNet(model =torchvision.models.resnext50_32x4d(pretrained=True),num_classes=num_classes,in_channel=in_channel)
    elif modelname == 'resnext101_32x8d':
        model = ResNet(model =torchvision.models.resnext101_32x8d(pretrained=True),num_classes=num_classes,in_channel=in_channel) 
    elif modelname == 'se_resnet34':
        model = ResNet(model=se_resnet.se_resnet34(),num_classes=num_classes,in_channel=in_channel)
    elif modelname == 'se_resnet50':
        model = ResNet(model=se_resnet.se_resnet50(pretrained=True),num_classes=num_classes,in_channel=in_channel)
    elif modelname == 'se_resnet_test':
        model = ResNet(model=se_resnet.se_resnet_test(),num_classes=num_classes,in_channel=in_channel)
    elif modelname == 'se_resnet101':
        model = ResNet(model=se_resnet.se_resnet101(),num_classes=num_classes,in_channel=in_channel)
    elif modelname == 'se_resnet152':
        model = ResNet(model=se_resnet.se_resnet152(),num_classes=num_classes,in_channel=in_channel)
    elif modelname == 'se_inception_v3':
        model = se_inception_v3()
    elif modelname[:12] == 'efficientnet':
        model = Efficientnet(model_name=modelname,num_classes=num_classes,in_channel=in_channel)
    elif modelname == 'EfficientNetv2':
        model = efficientnetv2_s(num_classes=num_classes,in_channel=in_channel)
    elif modelname == 'vit':
        efficient_transformer = Linformer(
        dim=128,
        seq_len=64+1,  # 8x8 patches + 1 cls-token
        depth=6,
        heads=4,
        k=64
        )
        model = ViT(
            dim=128,
            image_size=128,
            patch_size=16,
            num_classes=num_classes,
            transformer=efficient_transformer,
            channels=3,
        )

    return model


def get_double_model(modelname,ratio,num_classes=2):
    if modelname =='resnet18':
        model1 = resnet18(pretrained=True)
        model2 = resnet18(pretrained=True)
        model = ResNet_double(model1 =model1,model2 =model2,ratio=[ratio,1-ratio],num_classes=num_classes)            
    elif modelname =='resnet34':
        model = ResNet_double(model1 =torchvision.models.resnet34(pretrained=True),model2 =torchvision.models.resnet34(pretrained=True),num_classes=num_classes)
    elif modelname =='resnet50':
        model1 = resnet50(pretrained=True,num_classes=num_classes)
        model2 = resnet50(pretrained=True,num_classes=num_classes)
        model = ResNet_double(model1 =model1,model2 =model2,ratio=[ratio,1-ratio],num_classes=num_classes)
        # model = ResNet(model1 =torchvision.models.resnet50(pretrained=True),model2 =torchvision.models.resnet50(pretrained=True))
    elif modelname =='resnext50_32x4d':
        model = ResNet_double(model1 =torchvision.models.resnext50_32x4d(pretrained=True),model2 =torchvision.models.resnext50_32x4d(pretrained=True),num_classes=num_classes )
    elif modelname == 'resnext101_32x8d':
        model = ResNet_double(model1 =torchvision.models.resnext101_32x8d(pretrained=True),model2 =torchvision.models.resnext101_32x8d(pretrained=True),num_classes=num_classes )
    elif modelname == 'se_resnet18':
        model1 = se_resnet18()
        model2 = se_resnet18()
        model = ResNet_double(model1 =model1,model2 =model2,ratio=[ratio,1-ratio],num_classes=num_classes)
    elif modelname == 'se_resnet50':
        # model = ResNet(model1=se_resnet50(pretrained=True),model2=se_resnet18())
        model1 = se_resnet50(pretrained=True)
        model2 = se_resnet50(pretrained=True)
        model = ResNet_double(model1 =model1,model2 =model2,ratio=[ratio,1-ratio],num_classes=num_classes)
    return model

def get_3d_model(modelname,num_classes=3,n_input_channels=2):
    if modelname[:6] =='resnet':
        model = resnet.generate_model(
                                n_classes=num_classes,
                                n_input_channels=n_input_channels,
                                model_depth = int(modelname[6:]))
    elif modelname[:5]=='senet':
        model = senet.se_resnet50(num_classes=num_classes,pretrained=None)
    return model